# -*- coding: utf-8 -*-
import os
import sys
import json
import argparse
from utils import Utils

class CodeBase:
	PROJ_DIR = "fuzz_projs"
	TAG_BUG = "mbug"
	TAG_FIX = "mfix"
	FUZZ_BRANCH = "mbug"
	DEFUALT_TARGET = "common/src/lib.rs"
	TAGET_REGEX = "pub fn fuzz_(\w+)\("

	def __init__(self):
		self.projs = {}
		self.tools = {}
		self.fuzz_targets = {}
		self.normal_targets = []
		self.check_update()
	
	def __get_repo_config(self):
		with open("projects.json", "r") as fp:
			return json.load(fp)

	def __load_defualt_fuzz_targets(self):
		import re
		with open(self.DEFUALT_TARGET, "r") as fp:
			data = fp.read()
		self.normal_targets = re.findall(self.TAGET_REGEX, data)

	def check_update(self):
		Utils.check_and_mkdir(self.proj_path)
		code_config = self.__get_repo_config()
		for repo in code_config:
			res = self.fetch_code(repo)
			if res:
				tag, uri = res
				self.update_proj(tag, uri=uri)
			else:
				Utils.err_print("Warning: Failed to fetch project " + repo["tag"] + " Url:" + repo["url"])
		self.__load_defualt_fuzz_targets()

	def fetch_code(self, repo):
		cmd_clone = "git clone {} {}"
		cmd_checkout = "git checkout -b {} {}"
		proj_dir = os.path.join(self.proj_path, repo["tag"])
		if os.path.exists(proj_dir):
			return repo["tag"], repo["url"]

		# step 1: git clone
		if not Utils.exec_command(cmd_clone.format(repo["url"], repo["tag"]), self.proj_path):
			return
		
		# step 2: create branches and set to bug branch
		res = False
		if repo["fix"]: # test the newest version, do not have fix repo
			res = Utils.exec_command(cmd_checkout.format(self.TAG_FIX, repo["fix"]), proj_dir)
		res |= Utils.exec_command(cmd_checkout.format(self.TAG_BUG, repo["bug"]), proj_dir)
		if not res:
			return
		return repo["tag"], repo["url"]

	def update_proj(self, tag, **kwargs):
		if tag not in self.projs:
			self.projs[tag] = {
				"uri": "",
				"fuzzer": {},
			}
		proj = self.projs[tag]
		if kwargs.get("uri"):
			proj["uri"] = kwargs.get("uri")
		if kwargs.get("fuzzer"):
			proj["fuzz"].update(kwargs.get("fuzzer"))

	def find_proj(self, tag):
		finds = []
		for proj in self.projs:
			if proj.startswith(tag):
				finds.append(proj)

		if len(finds) > 1:
			Utils.notice_print(f"There are {len(finds)} projects starts with {tag}.")
			return None
		return None if len(finds) == 0 else finds[0]

	@property
	def proj_path(self):
		return os.path.join(os.getcwd(), self.PROJ_DIR)

	def list_projects(self):
		Utils.notice_print("Current Fuzzing Projects:")

		for tag, item in self.projs.items():
			url = item["uri"]
			fuzzers = item["fuzzer"]
			Utils.notice_print(f"\t{tag} -> {url}")
			supported_fuzzers = []
			for key in self.fuzz_targets:
				flat_tag = tag.replace("-", "_")
				if flat_tag in key:
					supported_fuzzers.append(key)
			tags = " | ".join(supported_fuzzers)
			Utils.notice_print(f"\t\tspecial fuzzing targets: [{tags}]")
		self.__show_normal_targets()

	def __show_normal_targets(self):
		tools = " | ".join(self.tools.keys())
		Utils.notice_print(f"Supported fuzzing tools: [{tools}]")
		targets = " | ".join(self.normal_targets)
		Utils.notice_print(f"Normal fuzzing targets: [{targets}]")

	def add_fuzzer(self, fuzzer):
		self.tools[fuzzer.TAG] = fuzzer
		targets = fuzzer.load_fuzzing_targets(self.projs.keys())
		self.fuzz_targets.update(targets)

	def run_fuzzing(self, fuzz_tag):
		target = self.fuzz_targets.get(fuzz_tag)
		if target:
			target["fuzzer"].run_fuzzing(target["tag"])
		else:
			Utils.err_print(f"These is no fuzz with tag {fuzz_tag}.")
			tags = " | ".join(self.fuzz_targets.keys())
			Utils.notice_print(f"Available tags: [{tags}]")

	def run_target(self, tool, target):
		if tool not in self.tools or target not in self.normal_targets:
			Utils.err_print(f"Tool [{tool}] or Target [{target}] is not supported.")
			self.__show_normal_targets()
			exit(1)
		fuzzer = self.tools[tool]
		fuzzer.setup_default_target(target)
		fuzzer.setup_default_seeds(target)
		fuzzer.run_fuzzing(FuzzProj.DEFAULT_PROJ)

class FuzzProj:
	TARGES = {"string-interner": "string_interner"} # tag map

	# need to override
	TAG = "fuzz"
	WORK_DIR = "afl_fuzz_targets"
	TARGET_DIR = "afl_fuzz_targets/src"
	CORPUS_DIR = "common/corpus"
	IN_DIR = ""
	CMD = "fuzz {}"
	DEFAULT_PROJ = "default"

	def __init__(self):
		pass

	def load_fuzzing_targets(self, projs):
		targets = {}
		all_rs_files = []
		for fi in os.listdir(self.TARGET_DIR):
			if fi.endswith(".rs"):
				all_rs_files.append(fi[:-3])
		# print(all_rs_files)
		for fi in all_rs_files:
			key = "{}_{}".format(self.TAG, fi)
			targets[key] = {
				"fuzzer": self,
				"tag": fi
			}
		return targets

	def run_fuzzing(self, proj):
		cmd = self.CMD.format(proj)
		Utils.exec_command(cmd, self.WORK_DIR)

	def setup_default_target(self, target):
		tpl_path = os.path.join(self.WORK_DIR, "template.rs")
		default_path = os.path.join(self.TARGET_DIR, "default.rs")
		with open(tpl_path, "r") as fp:
			code = fp.read()
		code = code.replace("####", f"fuzz_{target}")
		with open(default_path, "w") as fp:
			fp.write(code)

	def setup_default_seeds(self, target):
		seed_path = os.path.join(self.CORPUS_DIR, target)
		in_path = os.path.join(self.WORK_DIR, self.IN_DIR)
		if os.path.exists(seed_path):
			Utils.remove_dir(in_path)
			# copy files
			Utils.copy_dir(seed_path, in_path)

class AFLFuzzer(FuzzProj):
	TAG = "afl"
	WORK_DIR = "afl_fuzz_targets"
	TARGET_DIR = "afl_fuzz_targets/src"
	IN_DIR = "in"
	CMD = "cargo afl fuzz -i {} -o out target/debug/{}"

	def run_fuzzing(self, proj):
		bin_path = os.path.join(self.WORK_DIR, "target/debug/{}".format(proj))
		
		if not os.path.isfile(bin_path):
			cmd = "cargo afl build --bin {}".format(proj)
			Utils.exec_command(cmd, self.WORK_DIR)

		in_path = os.path.join(self.WORK_DIR, self.IN_DIR)
		if not os.path.exists(in_path):
			# set a random input file to make afl run
			in_path = "src"

		out_path = os.path.join(self.WORK_DIR, "out")
		# Utils.remove_dir(out_path)
		cmd = self.CMD.format(self.IN_DIR, proj)
		Utils.exec_command(cmd, self.WORK_DIR)

class CargoFuzzer(FuzzProj):
	TAG = "cargo"
	WORK_DIR = "cargo_fuzz_targets"
	TARGET_DIR = "cargo_fuzz_targets/fuzz/fuzz_targets"
	IN_DIR = "fuzz/corpus/default"
	CMD = "cargo +nightly fuzz run {}"

class HfuzzFuzzer(FuzzProj):
	TAG = "hfuzz"
	WORK_DIR = "hfuzz_targets"
	IN_DIR = "hfuzz_workspace/default/input"
	TARGET_DIR = "hfuzz_targets/src"
	# FALGS = 'RUSTFLAGS="-Z sanitizer=address"'
	CMD = 'RUSTFLAGS="-Z sanitizer=address" cargo hfuzz run {}'

class FuzzchckFuzzer(FuzzProj):
	TAG = "fzcheck"
	WORK_DIR = "fzcheck_targets"
	IN_DIR = "fuzz/corpora/default"
	TARGET_DIR = "fzcheck_targets/fuzz/instrumented/src"
	# FALGS = 'RUSTFLAGS="-Z sanitizer=address"'
	CMD = 'cargo fuzzcheck run {} fuzz'

	def run_fuzzing(self, proj):
		# clean compiled project
		bin_dir = os.path.join(self.WORK_DIR, "fuzz/instrumented/target")
		Utils.remove_dir(bin_dir)
		super(FuzzchckFuzzer, self).run_fuzzing(proj)

def setup_args():
	parser = argparse.ArgumentParser(prog="fuzz_tool", epilog="e.g. python3 code_setup.py -t afl base64")
	parser.add_argument("-a", "--all", action="store_true", help="Fuzzing all projects.")
	parser.add_argument("-l", "--list", action="store_true", help="List current projects.")
	parser.add_argument("-f", "--fuzz", metavar="target", help="Choose fuzz target to perform fuzzing.")
	parser.add_argument("-t", "--target", nargs=2, metavar=("tool", "target"), help="Choose a fuzzing tool and target.")
	args = parser.parse_args()
	if len(sys.argv) == 1:
	# if not any(vars(args).values()):
		parser.print_help()
		sys.exit(1)
	return args

def main(args):
	# HfuzzFuzzer().run_fuzzing("guess")
	fuzz_projs = CodeBase()
	fuzz_projs.add_fuzzer(AFLFuzzer())
	fuzz_projs.add_fuzzer(CargoFuzzer())
	fuzz_projs.add_fuzzer(HfuzzFuzzer())
	fuzz_projs.add_fuzzer(FuzzchckFuzzer())
	# print(args)
	if args.all:
		pass
	elif args.list:
		fuzz_projs.list_projects()
	elif args.fuzz:
		fuzz_projs.run_fuzzing(args.fuzz)
	elif args.target:
		tool, target = args.target
		fuzz_projs.run_target(tool, target)

if __name__ == "__main__":
	args = setup_args()
	main(args)
